import asyncio
import argparse
import random
from datetime import datetime
from src.config.config_loader import ConfigLoader
from src.data.question_loader import QuestionLoader
from src.benchmark.benchmark import Benchmark
from src.utils.monitor import SystemMonitor

async def main():
    # 打印当前模型测试摘要
    parser = argparse.ArgumentParser(description="vLLM模型性能测试")
    parser.add_argument("--concurrency", type=str, default="1", 
                        help="要测试的并发级别，逗号分隔")
    args = parser.parse_args()
    
    # 初始化组件
    config_loader = ConfigLoader()
    question_loader = QuestionLoader()
    benchmark = Benchmark(config_loader)
    
    # 加载问题并解析并发级别
    questions = question_loader.load_questions()
    concurrency_levels = [int(x) for x in args.concurrency.split(",")]
    
    # 打印测试摘要
    print(f"vLLM {config_loader.model_name}模型性能测试")
    print(f"{'=' * 50}")
    print(f"测试开始时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"测试并发级别: {concurrency_levels}")
    print(f"服务器地址: {config_loader.server_url}")
    print(f"{'=' * 50}\n")
    # 调用资源监控
    monitor = SystemMonitor()
    monitor.start()
    
    # 执行基准测试
    all_results = []
    for concurrency in concurrency_levels:
        # 获取问题库随机问题，并执行基准测试
        selected_questions = random.choices(questions, k=concurrency)
        result = await benchmark.run_benchmark(concurrency, selected_questions)
        if result:
            all_results.append(result)
    
    # 打印所有基准测试结果
    if all_results:
        print("\n所有基准测试结果摘要:")
        print("=" * 115)
        
        header_format = "{:<12} {:<16} {:<18} {:<17} {:<18} {:<18}"
        data_format = "{:<15} {:<20.2f} {:<20} {:<20.2f} {:<24.2f} {:<20.2f}"
        
        print(header_format.format(
            "并发数",
            "总耗时(秒)",
            "总Token数",
            "每秒Token数",
            "每线程每秒Token数",
            "平均Token数"
        ))
        print("-" * 115)
        
        for result in all_results:
            print(data_format.format(
                result['concurrency'],
                result['total_time'],
                result['total_tokens'],
                result['tokens_per_second'],
                result['tokens_per_thread'],
                result['avg_tokens_per_request']
            ))
        print("=" * 115)


    # 获取监控结果
    stats = monitor.end()
    print(f"\n{' 性能摘要 '.center(50, '=')}")
    print(f"平均CPU使用率: {stats['avg_cpu']:.2f}%")
    print(f"平均内存使用率: {stats['avg_mem']:.2f}%")
    print('=' * 50)
    print(f"测试结束时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    # 调用GPU监控  
    # gpu_util, mem_used = monitor.collect_gpu_metrics()
    # if gpu_util is None:
    #     print("[硬件监控] 采集失败")
    #     return
    # else:
    #     print(f"[硬件监控] 采集成功 GPU利用率={gpu_util}% 显存使用={mem_used}MB")



if __name__ == "__main__":
    asyncio.run(main())